import cv2
import numpy as np
from moviepy.video.io.VideoFileClip import VideoFileClip


def save_mask_as_video(mask_sequence, out_path, fps=30, frame_size=None):
    if frame_size is None:
        h, w = mask_sequence[0].shape
        frame_size = (w, h)

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    video_save = cv2.VideoWriter(out_path, fourcc, fps, frame_size, isColor=False)

    for mask in mask_sequence:
        mask_uint8 = (mask * 255).astype(np.uint8) if mask.max() <= 1.1 else mask.astype(np.uint8)
        video_save.write(mask_uint8)

    video_save.release()
    print(f"Mask video saved to {out_path}")


def save_fg_as_video(frame_fg_sequence, out_path, fps=30, frame_size=None):
    if frame_size is None:
        h, w = frame_fg_sequence[0].shape[:2]
        frame_size = (w, h)

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    video_save = cv2.VideoWriter(out_path, fourcc, fps, frame_size, isColor=True)

    for frame_fg in frame_fg_sequence:
        video_save.write(frame_fg.astype(np.uint8))

    video_save.release()
    print(f"Foreground video saved to {out_path}")


def read_video(matting_out_path, gray=True):
    cap = cv2.VideoCapture(matting_out_path)
    if not cap.isOpened():
        print("无法打开视频文件")
        return None

    frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if gray:
            frame = np.expand_dims(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY), axis=-1)
        frames.append(frame / 255.0)

    cap.release()
    return frames


def add_audio_to_video(audio_video_path, silent_video_path, output_path):
    # Use context managers to ensure resources are properly managed
    with VideoFileClip(audio_video_path) as original_video, VideoFileClip(silent_video_path) as silent_video:
        # Set the audio of the output video as the audio of the original video
        audio = original_video.audio
        final_video = silent_video.set_audio(audio)
        # Write the final video file with the specified codec
        final_video.write_videofile(output_path, codec="libx264", audio_codec="aac")
